from django.shortcuts import HttpResponse
from langchain_community.document_loaders import TextLoader,PyPDFLoader,CSVLoader,Docx2txtLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.document_loaders import TextLoader
from langchain_community.vectorstores import Neo4jVector
from langchain_core._api import LangChainDeprecationWarning
from model.my_chat_model import ChatModel
import json
import os
import warnings

warnings.filterwarnings("ignore", category=LangChainDeprecationWarning)

#文档存储
def doc_store(file_name):
    path = "../static/file"
    if file_name.endswith(".pdf"):
        loader = PyPDFLoader(f"{path}/{file_name}")
    elif file_name.endswith(".docx"):
        loader = Docx2txtLoader(f"{path}/{file_name}")
    elif file_name.endswith(".csv"):
        loader = CSVLoader(f"{path}/{file_name}",encoding="utf-8")
    else:
        loader = TextLoader(f"{path}/{file_name}",encoding="utf-8")
    text_splitter = RecursiveCharacterTextSplitter(
        # 优先按结构切分
        separators=[
            "\n\n",  # 段落之间
            "\n",  # 换行（句子之间）
            "。", "！", "？",  # 中文句号、感叹号、问号
            "；", "……", "…",  # 中文分号、省略号
            " ",  # 空格（词之间）
            ""  # 最后才按字符切
        ],
        chunk_size=300,  # 每块约 150-200 个汉字（合理上下文长度）
        chunk_overlap=60,  # 保留部分上下文
        length_function=len  # 可替换为 tokenizer 计算 token 数
    )
    splited_docs=loader.load_and_split(text_splitter=text_splitter)
    print(f"文档分割完毕，一共有{len(splited_docs)}")

    #获取模型
    chat=ChatModel()
    embedding_model=chat.get_embedding_model()
    #创建neo4j存储
    batch_size=10
    for i in range(0, len(splited_docs), batch_size):
        batch_docs=splited_docs[i:i+batch_size]
        vetor_store=Neo4jVector.from_documents(
            documents=batch_docs,
            embedding=embedding_model,
            url=os.getenv("NEO4J_URI"),
            username=os.getenv("NEO4J_USERNAME"),
            password=os.getenv("NEO4J_PASSWORD"),
            index_name="finance_chunks",
            node_label="FinanceChunks",
            text_node_property="text",
            embedding_node_property="embedding",
        )
        print(f"第{i}/{len(splited_docs)}批数据成功存入Neo4j数据库")
    print("所有数据已成功存入Neo4j数据库")

def upload_file(request):
    if request.method == "POST":
        file = request.FILES.get("file")
        print(f"file={file}")
        with open(f"./static/file/{file}", "wb") as f:
            for chunk in file.chunks():
                f.write(chunk)
        #文档入库
        doc_store(file.name)
    return HttpResponse(json.dumps({"code":200,"msg":"success"}))

